import os 
import logging
import datetime as dt
import polars as pl
import time

from binance.spot import Spot as Client
from binance.lib.utils import config_logging

from cryptoAlpha.dataFeed.IDataFeed import IDataFeed
from cryptoAlpha.tools.timeTool import timestamp_to_datetime 

class BinanceDataFeed(IDataFeed):
    def __init__(self, api_key: str = None, api_secret: str = None, is_testNet: bool = False, logLevel=logging.DEBUG) -> None:
        """
        is_testNet bool: 是否使用测试网络
        """
        super().__init__("BinanceDataFeed")
        config_logging(logging, logLevel)
        bese_url = "https://testnet.binance.vision" if is_testNet else "https://api.binance.com"
        self.exchange = "Binance"
        self.client = Client(
            api_key=api_key if api_key else os.environ.get("BINANCE_API_KEY"),
            api_secret=api_secret if api_secret else os.environ.get("BINANCE_API_SECRET"),
            base_url=bese_url)
        
    def get_exchange_info(self, symbol: str = None, symbols: list = None, permissions: list = None):
        result = self.client.exchange_info(symbol, symbols, permissions)
        
        return result
    
    def get_depth(self, symbol: str, limit: int=100):
        result = self.client.depth(symbol, limit=limit)

        return result
    
    def get_historical_trades(self, symbol: str, fromId: int=None, limit=1000):
        for i in range(10):
            try:
                result = self.client.historical_trades(symbol, fromId=fromId, limit=limit)
                if result is None or len(result) == 0:
                    return None
                    
                df = pl.DataFrame(result)

                df = df.with_columns(
                    trading_date = pl.col("time").map_elements(lambda x:timestamp_to_datetime(x).date(), pl.Date),
                    datetime = pl.col("time").map_elements(timestamp_to_datetime, pl.Datetime),
                    symbol = pl.lit(symbol),
                    exchange = pl.lit(self.exchange),
                )
                df = df[["trading_date", "datetime", "symbol", "exchange", "id", "price", "qty", "quoteQty", "isBuyerMaker", "isBestMatch"]]
                df = pl.DataFrame(df, schema= {
                    "trading_date": pl.Date,
                    "datetime": pl.Datetime,
                    "symbol": pl.Categorical,
                    "exchange": pl.Categorical,
                    "id": pl.UInt64,
                    "price": pl.Float64,
                    "qty": pl.Float64,
                    "quoteQty": pl.Float64,
                    "isBuyerMaker": pl.Boolean,
                    "isBestMatch": pl.Boolean,
                })
                if df is None:
                    return None
                else:
                    df = df.sort("datetime").unique(subset=["datetime"])
                    return df
            except Exception as e:
                print(e)
                time.sleep(10)
    
    def get_agg_trades(self, symbol: str, fromId:int=None, limit: int=1000, aggSpan: int=10000, sleep: int=0,):
        """
        接口中的agg_traders只能获取很短一段时间的聚合数据，这里通过获取原始的trade来自己聚合
        获取从fromId到最新的数据
        """
        def agg_trades(df):
            return df.group_by(["trading_date", "datetime", "symbol", "exchange", "price" ,"isBuyerMaker"]).agg(
                first_id = pl.col("id").first(),
                last_id = pl.col("id").last(),
                qty = pl.col("qty").sum(),
                qtyList = pl.col("qty"),
            )
        
        aggTrades = None
        fromId = fromId
        df = self.get_historical_trades(symbol, fromId=fromId, limit=limit)
        fromId = df["id"].max()
        while True:
            part = self.get_historical_trades(symbol, fromId=fromId, limit=limit)
            if part is None or len(part) == 0:
                break
            
            if df is None:
                df = part
            else:
                df = df.vstack(part)

            fromId = df["id"].max()
            
            if len(df) > aggSpan:
                aggDf = agg_trades(df)
                if aggTrades is None:
                    aggTrades = aggDf
                else:
                    aggTrades = aggTrades.vstack(aggDf)
                    
                df = None

            if sleep > 0:
                time.sleep(sleep)
        
        if aggTrades is None:
            return None
        else:
            # 最后一次合并
            aggTrades = aggTrades.sort("datetime")
            aggTrades = aggTrades.group_by(["trading_date","datetime","symbol","exchange","price","isBuyerMaker"]).agg(qtyList = pl.col("qtyList").flatten())
            return aggTrades
    
    def get_klines(self, symbol: str, interval: str, startTime: dt.datetime = None, endTime: dt.datetime = None):
        startTimeStamp = int(startTime.timestamp() * 1000) if startTime else None
        endTimeStamp = int(endTime.timestamp() * 1000) if endTime else None
        last = None
        df = None

        while True:
            result = self.client.klines(symbol, interval, startTime=startTimeStamp, endTime=endTimeStamp)
            if result is None or len(result) == 0:
                break

            for ts, open, high, low, close, volume, closeTs, turnover, tradeNum, bidVol, bidTurnover, _ in result:
                bar = pl.DataFrame({
                    "trading_date": timestamp_to_datetime(ts).date(),
                    "datetime" : timestamp_to_datetime(ts),
                    "symbol" : symbol,
                    "exchange" : self.exchange,
                    "open" : float(open),
                    "high" : float(high),
                    "low" : float(low),
                    "close" : float(close),
                    "volume" : float(volume),
                    "turnover" : float(turnover),
                    "tradeNum" : float(tradeNum),
                    "bidVol" : float(bidVol),
                    "bidTurnover" : float(bidTurnover),
                }, schema={
                    "trading_date" : pl.Date,
                    "datetime" : pl.Datetime,
                    "symbol" : pl.Categorical,
                    "exchange" : pl.Categorical,
                    "open" : pl.Float64,
                    "high" : pl.Float64,
                    "low" : pl.Float64,
                    "close" : pl.Float64,
                    "volume" : pl.Float64,
                    "turnover" : pl.Float64,
                    "tradeNum" : pl.Float64,
                    "bidVol" : pl.Float64,
                    "bidTurnover" : pl.Float64,
                })
                
                if df is None:
                    df = pl.DataFrame(bar)
                else:
                    df = df.vstack(bar)
            
            last = int(df["datetime"].max().timestamp() * 1000)
            if last >= endTime.timestamp() * 1000 or startTimeStamp==last:
                break
            startTimeStamp = last

        if df is None:
            return None
        else:
            df = df.filter((df["datetime"] >= startTime) & (df["datetime"] <= endTime))
            df = df.sort("datetime").unique(subset=["datetime"])
            return df
    
if __name__ == "__main__":
    feed = BinanceDataFeed(is_testNet=True)
    info = feed.get_exchange_info()
    
    startTime = dt.datetime(2024,9,6)
    endTime = dt.datetime(2024,9,7)
    
    def on_day(df):
        pass
    # r = feed.get_klines("BTCUSDT", "1m", startTime, endTime)
    # r = feed.get_agg_trades("BTCUSDT", startTime, endTime)
    r = feed.get_agg_trades("BTCUSDT", fromId=0, limit=1000, on_day=on_day)

    print(r)

 